const CodeWriter = @This();
const std = @import("std");
const Command = @import("Command.zig");
const String = []const u8;

writer: std.fs.File.Writer,

filename_buf: [std.fs.max_name_bytes]u8 = undefined,
filename: []u8 = undefined,
label_prefix_buf: [1024]u8 = undefined,
label_prefix: String = undefined, // "file.func" or "file" if func is null

const @"A=SP-1" =
    \\@SP
    \\A=M-1
;

const @"*SP=D, SP++" =
    \\@SP
    \\A=M
    \\M=D
    \\@SP
    \\M=M+1
;

const @"SP--, D=*SP" =
    \\@SP
    \\AM=M-1
    \\D=M
;

const @"SP--, D=*SP, A=SP-1" =
    \\@SP
    \\AM=M-1
    \\D=M
    \\A=A-1
;

const map = struct {
    const binary: std.StaticStringMap(String) = .initComptime(.{
        .{ "add", "D+M" },
        .{ "and", "D&M" },
        .{ "sub", "M-D" },
        .{ "or", "D|M" },
    });

    const unary: std.StaticStringMap(u8) = .initComptime(.{
        .{ "neg", '-' },
        .{ "not", '!' },
    });

    const segments: std.StaticStringMap(String) = .initComptime(.{
        .{ "local", "LCL" },
        .{ "argument", "ARG" },
        .{ "this", "THIS" },
        .{ "that", "THAT" },
    });
};

pub fn init(w: std.fs.File.Writer) CodeWriter {
    return .{
        .writer = w,
    };
}

pub fn startupCode(cw: CodeWriter) void {
    // write bootstrap code
    // SP=256
    cw.write(
        \\// bootstrap code
        \\@256
        \\D=A
        \\@SP
        \\M=D
        \\
    , .{});

    // call Sys.init
    cw.write("// call Sys.init\n", .{});
    cw.call("Sys.init", 0);
}

pub fn translate(cw: *CodeWriter, cmd: Command) !void {
    switch (cmd.type) {
        .push, .pop => cw.pushpop(cmd),
        .arithmetic => cw.arithmetic(cmd.arg1),
        .label => cw.label(cmd.arg1),
        .goto => cw.goto(cmd.arg1),
        .@"if-goto" => cw.if_goto(cmd.arg1),
        .call => cw.call(cmd.arg1, cmd.arg2.?),
        .function => try cw.function(cmd.arg1, cmd.arg2.?),
        .@"return" => try cw.@"return"(),
    }
}

pub fn setFileName(cw: *CodeWriter, file: String) !void {
    cw.filename = try std.fmt.bufPrint(&cw.filename_buf, "{s}", .{file});
    try cw.setPrefix(cw.filename);
}

fn setPrefix(cw: *CodeWriter, prefix: String) !void {
    cw.label_prefix = try std.fmt.bufPrint(&cw.label_prefix_buf, "{s}", .{prefix});
}

fn write(cw: CodeWriter, comptime format: String, args: anytype) void {
    cw.writer.print(format, args) catch unreachable;
}

fn arithmetic(cw: CodeWriter, op: String) void {
    const S = struct {
        var index: usize = 0; // static local variable
    };

    switch (op[0]) {
        'e', 'g', 'l' => { // eq, gt, lt
            var output: [2]u8 = undefined;
            // 生成的 label 前加 '$'，防止和 vm 中的 label 冲突
            cw.write(
                \\{0s}
                \\D=M-D
                \\@${1s}_TRUE{2d}
                \\D;J{1s}
                \\{3s}
                \\M=0
                \\@${1s}_END{2d}
                \\0;JMP
                \\(${1s}_TRUE{2d})
                \\{3s}
                \\M=-1
                \\(${1s}_END{2d})
                \\
            , .{
                @"SP--, D=*SP, A=SP-1",
                std.ascii.upperString(&output, op),
                S.index,
                @"A=SP-1",
            });
            S.index += 1;
        },
        'a', 's', 'o' => { // add, and, sub, or
            cw.write(
                \\{s}
                \\M={s}
                \\
            , .{
                @"SP--, D=*SP, A=SP-1",
                map.binary.get(op).?,
            });
        },
        'n' => { // neg, not
            cw.write(
                \\{s}
                \\M={c}M
                \\
            , .{
                @"A=SP-1",
                map.unary.get(op).?,
            });
        },
        else => unreachable,
    }
}

fn pushpop(cw: CodeWriter, cmd: Command) void {
    const segment = cmd.arg1;
    const index = cmd.arg2.?;

    switch (segment[0]) {
        'c' => { // constant
            cw.write(
                \\@{d}
                \\D=A
                \\{s}
                \\
            , .{
                index,
                @"*SP=D, SP++",
            });
        },
        's' => { // static
            if (cmd.type == .push) {
                cw.write(
                    \\@{s}.{d}
                    \\D=M
                    \\{s}
                    \\
                , .{
                    cw.filename,
                    index,
                    @"*SP=D, SP++",
                });
            } else {
                cw.write(
                    \\{s}
                    \\@{s}.{d}
                    \\M=D
                    \\
                , .{
                    @"SP--, D=*SP",
                    cw.filename,
                    index,
                });
            }
        },
        'p' => { // pointer
            const address = if (index == 0) "THIS" else "THAT";
            if (cmd.type == .push) {
                cw.write(
                    \\@{s}
                    \\D=M
                    \\{s}
                    \\
                , .{
                    address,
                    @"*SP=D, SP++",
                });
            } else {
                cw.write(
                    \\{s}
                    \\@{s}
                    \\M=D
                    \\
                , .{
                    @"SP--, D=*SP",
                    address,
                });
            }
        },
        't', 'l', 'a' => {
            if (std.mem.eql(u8, segment, "temp")) {
                if (cmd.type == .push) {
                    cw.write(
                        \\@{d}
                        \\D=M
                        \\{s}
                        \\
                    , .{
                        5 + index,
                        @"*SP=D, SP++",
                    });
                } else {
                    cw.write(
                        \\{s}
                        \\@{d}
                        \\M=D
                        \\
                    , .{
                        @"SP--, D=*SP",
                        5 + index,
                    });
                }
            } else { // local, argument, this, that
                if (cmd.type == .push) {
                    cw.write(
                        \\@{s}
                        \\D=M
                        \\@{d}
                        \\A=D+A
                        \\D=M
                        \\{s}
                        \\
                    , .{
                        map.segments.get(segment).?,
                        index,
                        @"*SP=D, SP++",
                    });
                } else {
                    cw.write(
                        \\@{s}
                        \\D=M
                        \\@{d}
                        \\D=D+A
                        \\@R13
                        \\M=D
                        \\{s}
                        \\@R13
                        \\A=M
                        \\M=D
                        \\
                    , .{
                        map.segments.get(segment).?,
                        index,
                        @"SP--, D=*SP",
                    });
                }
            }
        },
        else => unreachable,
    }
}

fn label(cw: CodeWriter, name: String) void {
    cw.write(
        \\({s}${s})
        \\
    , .{ cw.label_prefix, name });
}

fn goto(cw: CodeWriter, name: String) void {
    cw.write(
        \\@{s}${s}
        \\0;JMP
        \\
    , .{ cw.label_prefix, name });
}

fn if_goto(cw: CodeWriter, name: String) void {
    cw.write(
        \\{s}
        \\@{s}${s}
        \\D;JNE
        \\
    , .{
        @"SP--, D=*SP",
        cw.label_prefix,
        name,
    });
}

fn call(cw: CodeWriter, name: String, n_args: usize) void {
    const S = struct {
        var index: usize = 0; // static local variable
    };

    // push returnAddress
    cw.write(
        \\@RA_CALL_{s}{d}
        \\D=A
        \\{s}
        \\
    , .{
        name,
        S.index,
        @"*SP=D, SP++",
    });

    // push LCL, ARG, THIS, THAT
    const symbols = [_]String{ "LCL", "ARG", "THIS", "THAT" };
    for (symbols) |symbol| {
        cw.write(
            \\@{s}
            \\D=M
            \\{s}
            \\
        , .{
            symbol,
            @"*SP=D, SP++",
        });
    }

    // LCL=SP, ARG=SP-5-nArgs, goto f
    cw.write(
        \\@SP
        \\D=M
        \\@LCL
        \\M=D
        \\@5
        \\D=D-A
        \\@{0d}
        \\D=D-A
        \\@ARG
        \\M=D
        \\@{1s}
        \\0;JMP
        \\(RA_CALL_{1s}{2d})
        \\
    , .{
        n_args,
        name,
        S.index,
    });

    S.index += 1;
}

fn function(cw: *CodeWriter, name: String, n_args: usize) !void {
    const S = struct {
        var index: usize = 0; // static local variable
    };

    cw.write("({s})\n", .{name});
    if (n_args == 0) return;

    // 多于一个参数，就循环执行
    if (n_args > 1) {
        cw.write(
            \\@{d}
            \\D=A
            \\@R14
            \\M=D
            \\($INIT_LCLVAR{d})
            \\
        , .{
            n_args,
            S.index,
        });
    }

    // 当前的 SP 指向 LCL，所以 push constant 0 相当于清空 local 段
    cw.pushpop(.{
        .type = .push,
        .arg1 = "constant",
        .arg2 = 0,
    });

    if (n_args > 1) {
        // FIX 书上写 DM 可以作为 dest，但是 cpu Emulator 只认 MD
        cw.write(
            \\@R14
            \\MD=M-1
            \\@$INIT_LCLVAR{d}
            \\D;JGT
            \\
        , .{S.index});
    }

    S.index += 1;
    try cw.setPrefix(name);
}

fn @"return"(cw: *CodeWriter) !void {
    // 0 个参数时, ARG 和 Return IP 是同一个，retAddr 会被覆盖，所以要先存起来
    // R15=LCL, R14=*(LCL-5)
    cw.write(
        \\@LCL
        \\D=M
        \\@R15
        \\M=D
        \\@5
        \\A=D-A
        \\D=M
        \\@R14
        \\M=D
        \\
    , .{});

    // *ARG=pop(), SP=ARG+1
    cw.write(
        \\{s}
        \\@ARG
        \\A=M
        \\M=D
        \\@ARG
        \\D=M+1
        \\@SP
        \\M=D
        \\
    , .{@"SP--, D=*SP"});

    // symbol = *(R15--)
    const symbols = [_]String{ "THAT", "THIS", "ARG", "LCL" };
    for (symbols) |symbol| {
        cw.write(
            \\@R15
            \\AM=M-1
            \\D=M
            \\@{s}
            \\M=D
            \\
        , .{symbol});
    }

    // goto *R14
    cw.write(
        \\@R14
        \\A=M
        \\0;JMP
        \\
    , .{});

    try cw.setPrefix(cw.filename);
}
